"""
Finite element operators module for DLR-FEM implementation.

This module implements the assembly of mass matrices (M_x, M_y) and 
stiffness matrices (A_x, A_y) for mass-lumped finite element methods.
"""

import numpy as np
from scipy import sparse
from typing import Tuple, Optional
from .mesh import RectangularMesh
from .basis import TensorProductBasis, LagrangeBasis


class FEMOperators:
    """
    Finite element operators for 2D problems using tensor product basis.
    
    This class assembles the mass matrices M_x, M_y and stiffness matrices A_x, A_y
    as defined in equations (2.11) and (2.12) of the DLR-FEM paper.
    """
    
    def __init__(self, mesh: RectangularMesh, basis: TensorProductBasis):
        """
        Initialize FEM operators.
        
        Args:
            mesh: Rectangular mesh
            basis: Tensor product basis functions
        """
        self.mesh = mesh
        self.basis = basis
        
        # Pre-assemble operators
        self._assemble_operators()
    
    def _assemble_operators(self):
        """Assemble all finite element operators."""
        print("Assembling finite element operators...")
        
        # For mass-lumped methods, we use diagonal mass matrices
        self.M_x = self._assemble_mass_matrix_1d('x')
        self.M_y = self._assemble_mass_matrix_1d('y') 
        
        # Stiffness matrices
        self.A_x = self._assemble_stiffness_matrix_1d('x')
        self.A_y = self._assemble_stiffness_matrix_1d('y')
        
        # Linear operators L_x = -M_x^{-1} A_x, L_y = -M_y^{-1} A_y
        self.L_x = self._compute_linear_operator(self.M_x, self.A_x)
        self.L_y = self._compute_linear_operator(self.M_y, self.A_y)
        
        print("Operators assembled successfully.")
    
    def _assemble_mass_matrix_1d(self, direction: str) -> np.ndarray:
        """
        Assemble 1D mass matrix using mass lumping.
        
        For mass lumping, we use diagonal matrices with Gauss-Lobatto weights.
        
        Args:
            direction: 'x' or 'y'
            
        Returns:
            Diagonal mass matrix
        """
        if direction == 'x':
            n_points = self.mesh.nx
            spacing = self.mesh.hx
            weights = self.basis.basis_x.weights
        elif direction == 'y':
            n_points = self.mesh.ny  
            spacing = self.mesh.hy
            weights = self.basis.basis_y.weights
        else:
            raise ValueError("Direction must be 'x' or 'y'")
        
        # Mass lumping: diagonal entries are the quadrature weights times grid spacing
        # The factor of spacing/2 comes from mapping from reference element [-1,1]
        mass_diag = np.zeros(n_points)
        
        # Interior nodes get full weight
        for i in range(n_points):
            if direction == 'x':
                mass_diag[i] = spacing / 2.0 * weights[i % len(weights)]
            else:
                mass_diag[i] = spacing / 2.0 * weights[i % len(weights)]
        
        # For uniform grids, we can simplify
        if len(weights) <= n_points:
            # Repeat weights pattern for multiple elements
            weight_pattern = np.tile(weights, n_points // len(weights) + 1)[:n_points]
            mass_diag = spacing / 2.0 * weight_pattern
        
        return np.diag(mass_diag)
    
    def _assemble_stiffness_matrix_1d(self, direction: str) -> np.ndarray:
        """
        Assemble 1D stiffness matrix.
        
        Args:
            direction: 'x' or 'y'
            
        Returns:
            Stiffness matrix
        """
        if direction == 'x':
            n_points = self.mesh.nx
            spacing = self.mesh.hx
            basis_1d = self.basis.basis_x
        elif direction == 'y':
            n_points = self.mesh.ny
            spacing = self.mesh.hy  
            basis_1d = self.basis.basis_y
        else:
            raise ValueError("Direction must be 'x' or 'y'")
        
        # For finite differences on uniform grid, we can use standard stencils
        A = np.zeros((n_points, n_points))
        
        # Use second-order finite difference stencil
        h = spacing
        
        # Interior points: -u_{i-1} + 2u_i - u_{i+1}
        for i in range(1, n_points - 1):
            A[i, i-1] = -1.0 / h**2
            A[i, i] = 2.0 / h**2  
            A[i, i+1] = -1.0 / h**2
        
        # Boundary conditions (homogeneous Neumann by default)
        # Left boundary: du/dx = 0 -> -u_0 + u_1 = 0
        A[0, 0] = 1.0 / h**2
        A[0, 1] = -1.0 / h**2
        
        # Right boundary: du/dx = 0 -> u_{n-1} - u_{n-2} = 0  
        A[-1, -2] = -1.0 / h**2
        A[-1, -1] = 1.0 / h**2
        
        return A
    
    def _compute_linear_operator(self, M: np.ndarray, A: np.ndarray) -> np.ndarray:
        """
        Compute linear operator L = -M^{-1}A.
        
        Args:
            M: Mass matrix
            A: Stiffness matrix
            
        Returns:
            Linear operator matrix
        """
        # Since M is diagonal, M^{-1} is easy to compute
        M_inv = np.diag(1.0 / np.diag(M))
        return -M_inv @ A
    
    def apply_linear_operator(self, W: np.ndarray, epsilon: float) -> np.ndarray:
        """
        Apply linear operator ε²(L_x W + W L_y^T) to matrix W.
        
        Args:
            W: Coefficient matrix (m × n)
            epsilon: Regularization parameter
            
        Returns:
            Result of linear operator application
        """
        term1 = self.L_x @ W  # L_x W
        term2 = W @ self.L_y.T  # W L_y^T
        return epsilon**2 * (term1 + term2)
    
    def apply_nonlinear_operator(self, W: np.ndarray) -> np.ndarray:
        """
        Apply nonlinear operator N(W) = W - W^{∘3} (Hadamard power).
        
        For Allen-Cahn equation: N(W) = W - W^3 (element-wise)
        
        Args:
            W: Coefficient matrix
            
        Returns:
            Result of nonlinear operator application
        """
        return W - np.power(W, 3)  # Element-wise cubic
    
    def get_mass_matrices(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get mass matrices M_x and M_y."""
        return self.M_x, self.M_y
    
    def get_stiffness_matrices(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get stiffness matrices A_x and A_y."""
        return self.A_x, self.A_y
    
    def get_linear_operators(self) -> Tuple[np.ndarray, np.ndarray]:
        """Get linear operators L_x and L_y."""
        return self.L_x, self.L_y
    
    def matrix_exponential_action(self, A: np.ndarray, U: np.ndarray, 
                                dt: float) -> np.ndarray:
        """
        Compute exp(dt * A) * U efficiently.
        
        For the linear solve in DLR-FEM, we need exp(τ ε² L_x) U and exp(τ ε² L_y) V.
        
        Args:
            A: Matrix in the exponential
            U: Matrix to multiply
            dt: Time step
            
        Returns:
            exp(dt * A) * U
        """
        from scipy.linalg import expm
        
        # For small matrices, direct computation
        if A.shape[0] < 100:
            exp_A = expm(dt * A)
            return exp_A @ U
        else:
            # For larger matrices, could use Krylov methods
            # For now, use direct computation
            exp_A = expm(dt * A)
            return exp_A @ U
    
    def generalized_qr(self, U: np.ndarray, M: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Generalized QR decomposition with respect to mass matrix M.
        
        Computes U = Q R where Q^T M Q = I.
        
        Args:
            U: Input matrix to decompose
            M: Mass matrix (positive definite)
            
        Returns:
            (Q, R): Orthogonal factor Q and upper triangular factor R
        """
        # Cholesky decomposition of mass matrix
        try:
            L = np.linalg.cholesky(M)
        except np.linalg.LinAlgError:
            # If M is not positive definite, regularize slightly
            M_reg = M + 1e-12 * np.eye(M.shape[0])
            L = np.linalg.cholesky(M_reg)
        
        # Transform to standard inner product
        U_tilde = L @ U
        
        # Standard QR decomposition
        Q_tilde, R = np.linalg.qr(U_tilde)
        
        # Transform back
        L_inv = np.linalg.inv(L)
        Q = L_inv @ Q_tilde
        
        return Q, R


class InitialConditions:
    """Class for setting up initial conditions for the Allen-Cahn equation."""
    
    @staticmethod
    def kiss_bubble(X: np.ndarray, Y: np.ndarray, 
                   centers: list = None, 
                   radius: float = 0.15,
                   epsilon: float = 0.02) -> np.ndarray:
        """
        Kiss-bubble initial condition from Eq. (4.2) of the paper.
        
        u_0(x,y) = -∑_{i=1}^2 tanh((√((x-x_i)² + (y-y_i)²) - R) / (√2 ε)) + 1
        
        Args:
            X, Y: Coordinate meshes
            centers: List of bubble centers [(x1, y1), (x2, y2)]
            radius: Bubble radius R
            epsilon: Interface parameter
            
        Returns:
            Initial condition values
        """
        if centers is None:
            # Default configuration for symmetry breaking
            centers = [(0.3, 0.5), (0.7, 0.5)]
        
        u0 = np.ones_like(X)
        
        for x_center, y_center in centers:
            dist = np.sqrt((X - x_center)**2 + (Y - y_center)**2)
            u0 += -np.tanh((dist - radius) / (np.sqrt(2) * epsilon))
        
        return u0
    
    @staticmethod
    def random_perturbation(X: np.ndarray, Y: np.ndarray,
                           amplitude: float = 0.1,
                           seed: Optional[int] = None) -> np.ndarray:
        """
        Random perturbation initial condition.
        
        Args:
            X, Y: Coordinate meshes
            amplitude: Perturbation amplitude
            seed: Random seed for reproducibility
            
        Returns:
            Initial condition values
        """
        if seed is not None:
            np.random.seed(seed)
        
        return amplitude * np.random.randn(*X.shape)
    
    @staticmethod
    def sine_wave(X: np.ndarray, Y: np.ndarray,
                 kx: int = 2, ky: int = 2,
                 amplitude: float = 1.0) -> np.ndarray:
        """
        Sine wave initial condition for testing.
        
        Args:
            X, Y: Coordinate meshes
            kx, ky: Wave numbers in x and y directions
            amplitude: Wave amplitude
            
        Returns:
            Initial condition values
        """
        return amplitude * np.sin(kx * np.pi * X) * np.sin(ky * np.pi * Y)


if __name__ == "__main__":
    # Test FEM operators
    from .mesh import create_uniform_mesh
    from .basis import create_basis_functions
    
    print("Testing FEM operators...")
    
    # Create mesh and basis
    mesh = create_uniform_mesh((0, 1), (0, 1), 10, 10)
    basis = create_basis_functions(poly_order=1)
    
    # Create operators
    operators = FEMOperators(mesh, basis)
    
    # Test matrices
    M_x, M_y = operators.get_mass_matrices()
    A_x, A_y = operators.get_stiffness_matrices()
    
    print(f"Mass matrix M_x shape: {M_x.shape}")
    print(f"Mass matrix M_y shape: {M_y.shape}")
    print(f"Stiffness matrix A_x shape: {A_x.shape}")
    print(f"Stiffness matrix A_y shape: {A_y.shape}")
    
    # Test initial conditions
    X, Y = mesh.get_coordinates()
    u0_kiss = InitialConditions.kiss_bubble(X, Y)
    u0_sine = InitialConditions.sine_wave(X, Y)
    
    print(f"Kiss-bubble IC shape: {u0_kiss.shape}")
    print(f"Kiss-bubble IC range: [{np.min(u0_kiss):.3f}, {np.max(u0_kiss):.3f}]")
    print(f"Sine wave IC shape: {u0_sine.shape}")
    print(f"Sine wave IC range: [{np.min(u0_sine):.3f}, {np.max(u0_sine):.3f}]")